#!/usr/bin/env python3
# Ensure the above line is the very first line of your script
# and that the script has execute permissions: chmod +x your_script_name.py

import requests
import json
import os
import argparse
from dotenv import load_dotenv

# Load environment variables from .env file
# This will load variables like CF_API_KEY, CF_EMAIL, CF_ZONE_ID from your .env file
load_dotenv()

# --- Configuration (Loaded from .env or environment variables) ---
CF_API_KEY = os.environ.get("CF_API_KEY")
CF_EMAIL = os.environ.get("CF_EMAIL")
CF_ZONE_ID = os.environ.get("CF_ZONE_ID") # The Cloudflare Zone ID for your domain

# Cloudflare API base URL
CF_API_BASE_URL = "https://api.cloudflare.com/client/v4"

def get_cloudflare_headers():
    """
    Returns the standard headers required for Cloudflare API requests.
    Checks if API Key and Email are loaded.
    """
    if not all([CF_API_KEY, CF_EMAIL]):
        print("Error: Cloudflare API Key (CF_API_KEY) or Email (CF_EMAIL) is not set.")
        print("Please ensure these are defined in your .env file or as environment variables.")
        exit(1)

    return {
        "X-Auth-Email": CF_EMAIL,
        "X-Auth-Key": CF_API_KEY,
        "Content-Type": "application/json"
    }

def get_all_a_records(zone_id):
    """
    Fetches all A type DNS records for a given Cloudflare zone ID, handling pagination.
    """
    if not zone_id:
        print("Error: Cloudflare Zone ID is not provided.")
        return []

    url = f"{CF_API_BASE_URL}/zones/{zone_id}/dns_records"
    headers = get_cloudflare_headers()
    all_records = []
    page = 1
    per_page = 100 # Maximum allowed per_page by Cloudflare API

    print(f"Fetching all A records for Zone ID: {zone_id}...")

    while True:
        params = {"page": page, "per_page": per_page, "type": "A"} # Filter by A records
        try:
            response = requests.get(url, headers=headers, params=params)
            response.raise_for_status() # Raise an HTTPError for bad responses (4xx or 5xx)
            data = response.json()

            if data["success"]:
                records_on_page = data["result"]
                all_records.extend(records_on_page)

                total_pages = data["result_info"]["total_pages"]
                print(f"  Fetched page {page}/{total_pages} ({len(records_on_page)} records). Total A records collected so far: {len(all_records)}")

                if page < total_pages:
                    page += 1
                else:
                    break # No more pages
            else:
                print(f"Error fetching DNS records for page {page}: {data['errors']}")
                return []
        except requests.exceptions.RequestException as e:
            if hasattr(e, 'response') and e.response is not None:
                print(f"Cloudflare API Error Response: {e.response.text}")
            print(f"Network error fetching DNS records on page {page}: {e}")
            return []
        except Exception as e:
            print(f"An unexpected error occurred while fetching DNS records on page {page}: {e}")
            return []

    return all_records

def update_dns_record(zone_id, record_id, record_name, new_ip, proxied_status):
    """
    Updates the IP address of a specific DNS A record.
    Preserves the proxied status.
    """
    if not all([zone_id, record_id, record_name, new_ip is not None]):
        print("Error: Missing parameters for updating DNS record.")
        return False

    url = f"{CF_API_BASE_URL}/zones/{zone_id}/dns_records/{record_id}"
    headers = get_cloudflare_headers()
    payload = {
        "type": "A",
        "name": record_name,
        "content": new_ip,
        "ttl": 3600, # Time to Live in seconds (1 hour recommended for dynamic DNS)
        "proxied": proxied_status # Preserve existing proxied status
    }

    print(f"Attempting to update DNS record '{record_name}' (ID: {record_id}) to new IP: {new_ip}...")

    try:
        response = requests.put(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        data = response.json()

        if data["success"]:
            print(f"Successfully updated DNS record '{record_name}' to new IP: {new_ip}")
            return True
        else:
            print(f"Error updating DNS record '{record_name}': {data['errors']}")
            return False
    except requests.exceptions.RequestException as e:
        if hasattr(e, 'response') and e.response is not None:
            print(f"Cloudflare API Error Response: {e.response.text}")
        print(f"Network error updating DNS record '{record_name}': {e}")
        return False
    except Exception as e:
        print(f"An unexpected error occurred while updating DNS record '{record_name}': {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Update all Cloudflare DNS A records pointing to a specific source IP to a new target IP.")
    parser.add_argument("--source-ip", required=True, help="The current IP address of records you want to replace.")
    parser.add_argument("--target-ip", required=True, help="The new IP address to set for matching records.")
    parser.add_argument("--confirm", action="store_true", help="Automatically confirm all updates without asking for user input for each record.")

    args = parser.parse_args()

    # Verify essential environment variables/config are loaded
    if not all([CF_API_KEY, CF_EMAIL, CF_ZONE_ID]):
        print("Error: Please ensure CF_API_KEY, CF_EMAIL, and CF_ZONE_ID environment variables are set.")
        print("These should be in your .env file located in the same directory as this script, or set in your system environment.")
        print("\nExample .env content:")
        print("CF_API_KEY=\"your_key\"")
        print("CF_EMAIL=\"your_email\"")
        print("CF_ZONE_ID=\"your_zone_id\"")
        return

    print(f"Starting mass IP update for Zone ID: {CF_ZONE_ID}")
    print(f"  Source IP: {args.source_ip}")
    print(f"  Target IP: {args.target_ip}\n")

    all_a_records = get_all_a_records(CF_ZONE_ID)
    if not all_a_records:
        print("No A records found or an error occurred. Aborting.")
        return

    matching_records_found = 0
    updated_records_count = 0

    records_to_update = []

    for record in all_a_records:
        record_id = record.get('id')
        record_name = record.get('name')
        current_ip = record.get('content')
        proxied_status = record.get('proxied')

        if record_id and record_name and current_ip:
            if current_ip == args.source_ip:
                matching_records_found += 1
                if current_ip == args.target_ip:
                    print(f"Record '{record_name}' is already set to the target IP '{args.target_ip}'. Skipping.")
                else:
                    records_to_update.append({
                        "id": record_id,
                        "name": record_name,
                        "current_ip": current_ip,
                        "proxied": proxied_status
                    })

    if not records_to_update:
        print(f"No records found with source IP '{args.source_ip}' that need updating to '{args.target_ip}'.")
        return

    print(f"\n--- Found {len(records_to_update)} record(s) to update out of {matching_records_found} matching source IP '{args.source_ip}' ---")
    for rec in records_to_update:
        print(f"  - Record: {rec['name']} (Current IP: {rec['current_ip']}, Proxied: {'Yes' if rec['proxied'] else 'No'})")
    print("--------------------------------------------------------------------------")

    if not args.confirm:
        user_input = input("Do you want to proceed with updating these records? (y/N): ").strip().lower()
        if user_input != 'y':
            print("Mass update cancelled by user.")
            return

    print("\n--- Executing Updates ---")
    for rec in records_to_update:
        print(f"Processing record: {rec['name']}...")
        success = update_dns_record(CF_ZONE_ID, rec['id'], rec['name'], args.target_ip, rec['proxied'])
        if success:
            updated_records_count += 1

    print(f"\n--- Update Summary ---")
    print(f"Total records found with source IP '{args.source_ip}': {matching_records_found}")
    print(f"Records successfully updated to '{args.target_ip}': {updated_records_count}")
    print(f"Records skipped (already target IP): {matching_records_found - len(records_to_update)}")
    print(f"Records that failed to update: {len(records_to_update) - updated_records_count}")
    print(f"--------------------------")


if __name__ == "__main__":
    main()

# Example usage:
# python update-dns.py --source-ip "1.2.3.4" --target-ip "5.6.7.8"